{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "OsrwCDQ5cviC"
   },
   "outputs": [],
   "source": [
    "# @title 💿 Installation\n",
    "\n",
    "!uv pip install -q openpipe-art==0.3.11.post5 langchain-core tenacity \"mcp>=1.11.0\" \"gql<4\" aiohttp --no-cache-dir"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "D8b8kgQ69ZDM"
   },
   "source": [
    "<a name=\"configuration\"></a>\n",
    "\n",
    "### 🎯 Configuration - Edit These Settings\n",
    "\n",
    "Add an OpenRouter API key and Smithery MCP server url below.\n",
    "\n",
    "### **Smithery Instructions**\n",
    "\n",
    "Smithery hosts a variety of useful MCP servers. If you're not sure which to use, try the [Exa server](https://smithery.ai/server/exa), which allows your model to query data from across the web.\n",
    "\n",
    "To generate an authenticated Smithery MCP url, follow these steps:\n",
    "\n",
    "1. Sign up for a [Smithery](https://smithery.ai) account\n",
    "2. Navigate to an MCP [server](https://smithery.ai/server/exa)\n",
    "3. Get an authenticated url by clicking the orange <u>Get URL with keys instead</u> button on the right\n",
    "4. Set `SMITHERY_MCP_URL` to the generated url"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "so6r1_OG9en3"
   },
   "outputs": [],
   "source": [
    "# Required - Used for generating training inputs and RULER evaluation\n",
    "OPENROUTER_API_KEY = \"\"  # Put your OpenRouter key here\n",
    "\n",
    "# 🔌 Point to any Smithery-hosted MCP server (make sure you click \"Get URL with keys instead\", otherwise this will not work)\n",
    "SMITHERY_MCP_URL = \"\"\n",
    "\n",
    "# Optional - Enables metric logging\n",
    "WANDB_API_KEY = \"\"\n",
    "\n",
    "# Choose the base model to train\n",
    "BASE_MODEL = \"Qwen/Qwen2.5-3B-Instruct\"  # Options: \"Qwen/Qwen2.5-3B-Instruct\", \"Qwen/Qwen2.5-7B-Instruct\", etc."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "caZYLROd8xnV"
   },
   "source": [
    "To teach a model to use your MCP server, click _Runtime_ and press _Run all_. Make sure you've enabled a free Tesla T4 GPU and edit the [configuration](#configuration) cell below!\n",
    "\n",
    "<div class=\"align-center\">\n",
    "<a href=\"https://github.com/openpipe/art\"><img src=\"https://github.com/openpipe/art/raw/main/assets/ART_pill.png\" height=\"50\"></a>\n",
    "<a href=\"https://discord.gg/zbBHRUpwf4\"><img src=\"https://github.com/openpipe/art/raw/main/assets/Discord_pill.png\" height=\"50\"></a>\n",
    "<a href=\"https://art.openpipe.ai\"><img src=\"https://github.com/openpipe/art/raw/main/assets/Documentation_pill.png\" height=\"50\"></a>\n",
    "\n",
    "Questions? Join the Discord and ask away! For feature requests or to leave a star, visit our [GitHub](https://github.com/openpipe/art).\n",
    "\n",
    "</div>\n",
    "\n",
    "<a href=\"https://art.openpipe.ai/\"><img src=\"https://github.com/openpipe/art/raw/main/assets/Header_separator.png\" height=\"5\"></a>\n",
    "\n",
    "**MCP•RL: Teach you agent how to use any MCP server**\n",
    "\n",
    "This notebook shows how to train a Qwen 2.5 3B model to effectively use any MCP server. Simply provide an MCP server url and the notebook will:\n",
    "\n",
    "1. Query the server's tools\n",
    "2. Generate a set of input tasks that use those tools\n",
    "3. Train the model on those tasks using automatic RULER evaluation\n",
    "4. Test the trained model by giving it new tasks to complete\n",
    "\n",
    "RULER judges response quality purely from the agent's final output - no labeled data required!\n",
    "\n",
    "*Note: In this notebook we use [Smithery](https://smithery.ai/) servers to keep things simple, but the technique below applies to all MCP servers!*\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "I_AFDSOv_LrB"
   },
   "outputs": [],
   "source": [
    "# @title Advanced Settings\n",
    "\n",
    "# Model configuration\n",
    "MODEL_NAME = \"mcprl-3b-exa\"  # Name for your trained model\n",
    "PROJECT_NAME = \"mcp-rl\"  # Project name for tracking\n",
    "\n",
    "# Training configuration\n",
    "TRAINING_CONFIG = {\n",
    "    \"num_training_inputs\": 16,  # Number of training inputs to generate\n",
    "    \"groups_per_step\": 2,  # Inputs to process per training step\n",
    "    \"num_epochs\": 1,  # Number of times through all data\n",
    "    \"rollouts_per_group\": 4,  # Different responses per input (for RULER comparison)\n",
    "    \"learning_rate\": 1e-5,  # Learning rate\n",
    "    \"max_training_steps\": None,  # Maximum training steps (set to None for no limit)\n",
    "}\n",
    "\n",
    "MAX_TURNS = 10  # Maximum number of turns for the model to generate during one rollout\n",
    "\n",
    "NUM_TEST_INPUTS = 8  # Number of test inputs to generate\n",
    "RULER_MODEL = \"openrouter/openai/o4-mini\"  # Model for RULER evaluation\n",
    "INPUT_GENERATION_MODEL = \"openai/o4-mini\"\n",
    "\n",
    "# GPU configuration (for T4 — keep these as-is unless you have a reason to change them)\n",
    "MAX_SEQ_LENGTH = 16384  # Maximum sequence length\n",
    "GPU_MEMORY_UTILIZATION = 0.7  # GPU memory usage (0.0-1.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "PfXGRuhhd7hr"
   },
   "outputs": [],
   "source": [
    "# @title Debug utilities\n",
    "\n",
    "import json\n",
    "import time\n",
    "import traceback\n",
    "from typing import Any\n",
    "\n",
    "DEBUG_LOG = True  # flip to False to silence logs\n",
    "LOG_JSON_MAX = 2000  # cap large JSON prints\n",
    "\n",
    "\n",
    "def _ts() -> str:\n",
    "    return time.strftime(\"%H:%M:%S\")\n",
    "\n",
    "\n",
    "def log(msg: str, **kv):\n",
    "    if not DEBUG_LOG:\n",
    "        return\n",
    "    parts = [f\"[{_ts()}] {msg}\"]\n",
    "    if kv:\n",
    "        kv_str = \" \".join(f\"{k}={repr(v)}\" for k, v in kv.items())\n",
    "        parts.append(\"| \" + kv_str)\n",
    "    print(\" \".join(parts))\n",
    "\n",
    "\n",
    "def log_json(title: str, payload: Any, max_len: int = LOG_JSON_MAX):\n",
    "    if not DEBUG_LOG:\n",
    "        return\n",
    "    try:\n",
    "        s = json.dumps(payload, indent=2, default=str)\n",
    "    except Exception:\n",
    "        s = str(payload)\n",
    "    if len(s) > max_len:\n",
    "        s = s[:max_len] + \"\\n... (truncated)\"\n",
    "    print(f\"[{_ts()}] {title}:\\n{s}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "gxUn4E_IPjq8"
   },
   "outputs": [],
   "source": [
    "# @title 🔌 MCP helpers\n",
    "\n",
    "from contextlib import asynccontextmanager\n",
    "\n",
    "import mcp.types as types\n",
    "from mcp.client.session import ClientSession\n",
    "from mcp.client.streamable_http import streamablehttp_client\n",
    "\n",
    "if not SMITHERY_MCP_URL:\n",
    "    raise ValueError(\"SMITHERY_MCP_URL is empty. Set it in the Configuration cell.\")\n",
    "\n",
    "\n",
    "@asynccontextmanager\n",
    "async def mcp_session():\n",
    "    \"\"\"\n",
    "    Connects to the remote Smithery MCP server using the full URL that includes\n",
    "    your API key & profile. No OAuth provider is used.\n",
    "    \"\"\"\n",
    "    async with streamablehttp_client(SMITHERY_MCP_URL) as (read, write, _):\n",
    "        async with ClientSession(read, write) as session:\n",
    "            await session.initialize()\n",
    "            yield session\n",
    "\n",
    "\n",
    "async def list_tools_and_resources():\n",
    "    \"\"\"Return (tools_result, resources_result) from the remote Smithery server.\"\"\"\n",
    "    async with mcp_session() as session:\n",
    "        tools = await session.list_tools()\n",
    "        try:\n",
    "            resources = await session.list_resources()\n",
    "        except Exception:\n",
    "            # Some servers don't implement resources; keep interface stable\n",
    "            class _Empty:\n",
    "                resources = []\n",
    "\n",
    "            resources = _Empty()\n",
    "        return tools, resources\n",
    "\n",
    "\n",
    "async def call_mcp_tool(tool_name: str, arguments: dict):\n",
    "    \"\"\"Invoke a tool on the remote Smithery server and return the CallToolResult.\"\"\"\n",
    "    async with mcp_session() as session:\n",
    "        return await session.call_tool(tool_name, arguments)\n",
    "\n",
    "\n",
    "tools, resources = await list_tools_and_resources()\n",
    "print(\"Tools:\", [t.name for t in tools.tools])\n",
    "print(\n",
    "    \"Resources:\",\n",
    "    [getattr(r, \"uri\", None) for r in getattr(resources, \"resources\", []) or []],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "nEB1JGY6Pjq8"
   },
   "outputs": [],
   "source": [
    "# @title Let's generate our train and validation scenarios!\n",
    "\n",
    "import os\n",
    "import random\n",
    "from collections import Counter\n",
    "from typing import Any, Dict, List\n",
    "\n",
    "import openai\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "load_dotenv()\n",
    "\n",
    "\n",
    "# ---------- lightweight \"nice print\" helpers (no extra deps) ----------\n",
    "class _C:\n",
    "    RESET = \"\\x1b[0m\"\n",
    "    DIM = \"\\x1b[2m\"\n",
    "    BOLD = \"\\x1b[1m\"\n",
    "    ITAL = \"\\x1b[3m\"\n",
    "    GRAY = \"\\x1b[90m\"\n",
    "    BLUE = \"\\x1b[34m\"\n",
    "    CYAN = \"\\x1b[36m\"\n",
    "    GREEN = \"\\x1b[32m\"\n",
    "    YELLOW = \"\\x1b[33m\"\n",
    "    RED = \"\\x1b[31m\"\n",
    "    MAGENTA = \"\\x1b[35m\"\n",
    "\n",
    "\n",
    "def _ts():\n",
    "    return time.strftime(\"%H:%M:%S\")\n",
    "\n",
    "\n",
    "def info(msg):\n",
    "    print(f\"[{_ts()}] {_C.BLUE}INFO{_C.RESET}  {msg}\")\n",
    "\n",
    "\n",
    "def step(msg):\n",
    "    print(f\"[{_ts()}] {_C.CYAN}STEP{_C.RESET}  {msg}\")\n",
    "\n",
    "\n",
    "def ok(msg):\n",
    "    print(f\"[{_ts()}] {_C.GREEN}OK{_C.RESET}    {msg}\")\n",
    "\n",
    "\n",
    "def warn(msg):\n",
    "    print(f\"[{_ts()}] {_C.YELLOW}WARN{_C.RESET}  {msg}\")\n",
    "\n",
    "\n",
    "def err(msg):\n",
    "    print(f\"[{_ts()}] {_C.RED}ERR{_C.RESET}   {msg}\")\n",
    "\n",
    "\n",
    "def dim(msg):\n",
    "    print(f\"{_C.DIM}{msg}{_C.RESET}\")\n",
    "\n",
    "\n",
    "def preview_scenarios(scenarios, n=5):\n",
    "    n = min(n, len(scenarios))\n",
    "    for i in range(n):\n",
    "        s = scenarios[i]\n",
    "        dim(\n",
    "            f\"   {i + 1}. {s['task'][:120].strip()}{'…' if len(s['task']) > 120 else ''}  \"\n",
    "            f\"{_C.GRAY}(difficulty {s['difficulty']}/5){_C.RESET}\"\n",
    "        )\n",
    "\n",
    "\n",
    "# ---------- required env/key check ----------\n",
    "# If OPENROUTER_API_KEY exists as a var, use it; otherwise pull from env\n",
    "_openrouter_key = os.getenv(\"OPENROUTER_API_KEY\")\n",
    "try:\n",
    "    _openrouter_key = _openrouter_key if _openrouter_key else OPENROUTER_API_KEY  # noqa: F821 (defined upstream in your notebook)\n",
    "except NameError:\n",
    "    pass\n",
    "\n",
    "if _openrouter_key:\n",
    "    os.environ[\"OPENROUTER_API_KEY\"] = _openrouter_key\n",
    "    ok(\"OPENROUTER_API_KEY found.\")\n",
    "else:\n",
    "    err(\"OPENROUTER_API_KEY is required for data generation and RULER evaluation.\")\n",
    "    raise ValueError(\n",
    "        \"OPENROUTER_API_KEY is required for data generation and RULER evaluation.\"\n",
    "    )\n",
    "\n",
    "\n",
    "# ---------- generator ----------\n",
    "async def generate_scenarios(\n",
    "    num_scenarios: int = 24,\n",
    ") -> List[Dict[str, Any]]:\n",
    "    t0 = time.perf_counter()\n",
    "    step(\"Fetching MCP tools & resources from remote server …\")\n",
    "    tools_result, resources_result = await list_tools_and_resources()\n",
    "    ok(f\"Fetched tools & resources in {time.perf_counter() - t0:.2f}s.\")\n",
    "\n",
    "    # summarize tools/resources\n",
    "    try:\n",
    "        tool_cnt = len(getattr(tools_result, \"tools\", []) or [])\n",
    "        res_cnt = len(getattr(resources_result, \"resources\", []) or [])\n",
    "    except Exception:\n",
    "        tool_cnt = res_cnt = 0\n",
    "    info(f\"Available: {tool_cnt} tool(s), {res_cnt} resource(s).\")\n",
    "\n",
    "    tools_info = []\n",
    "    for tool in tools_result.tools or []:\n",
    "        tools_info.append(\n",
    "            {\n",
    "                \"name\": tool.name,\n",
    "                \"description\": tool.description,\n",
    "                \"parameters\": tool.inputSchema,\n",
    "            }\n",
    "        )\n",
    "\n",
    "    resources_info = []\n",
    "    for resource in getattr(resources_result, \"resources\", []) or []:\n",
    "        resources_info.append(\n",
    "            {\n",
    "                \"uri\": str(resource.uri),\n",
    "                \"name\": resource.name,\n",
    "                \"description\": resource.description,\n",
    "                \"mimeType\": resource.mimeType,\n",
    "            }\n",
    "        )\n",
    "\n",
    "    step(\"Preparing prompt & JSON schema …\")\n",
    "    tools_description = json.dumps(tools_info, indent=2)\n",
    "    resources_description = (\n",
    "        json.dumps(resources_info, indent=2)\n",
    "        if resources_info\n",
    "        else \"No resources available\"\n",
    "    )\n",
    "\n",
    "    prompt = f\"\"\"You are an expert at creating realistic scenarios for testing AI agents that interact with MCP (Model Context Protocol) servers.\n",
    "\n",
    "Given the following available tools and resources from an MCP server, generate {num_scenarios} diverse, realistic scenarios that a user might want to accomplish using these tools.\n",
    "\n",
    "AVAILABLE TOOLS:\n",
    "{tools_description}\n",
    "\n",
    "AVAILABLE RESOURCES:\n",
    "{resources_description}\n",
    "\n",
    "Requirements for scenarios:\n",
    "1. Each scenario should be a task that can be accomplished using the available tools\n",
    "2. Scenarios should vary in complexity - some simple (1-2 tool calls), some complex (multiple tool calls)\n",
    "3. Scenarios should cover different use cases and tool combinations (though the task should not specify which tools to use)\n",
    "4. Each scenario should be realistic - something a real user might actually want to do\n",
    "5. Assign a difficulty rating from 1 (easy, single tool call) to 5 (hard, complex multi-step analysis)\n",
    "6. The task should always include generating a summary of the work done and a thorough analysis and report of the results\n",
    "\n",
    "You must respond with a JSON object containing a \"scenarios\" array of exactly {num_scenarios} objects. Each object must have:\n",
    "- \"task\": string describing the scenario\n",
    "- \"difficulty\": integer from 1-5 representing complexity\n",
    "\"\"\"\n",
    "\n",
    "    response_schema = {\n",
    "        \"type\": \"object\",\n",
    "        \"properties\": {\n",
    "            \"scenarios\": {\n",
    "                \"type\": \"array\",\n",
    "                \"items\": {\n",
    "                    \"type\": \"object\",\n",
    "                    \"properties\": {\n",
    "                        \"task\": {\"type\": \"string\"},\n",
    "                        \"difficulty\": {\"type\": \"integer\", \"minimum\": 1, \"maximum\": 5},\n",
    "                    },\n",
    "                    \"required\": [\"task\", \"difficulty\"],\n",
    "                    \"additionalProperties\": False,\n",
    "                },\n",
    "                \"minItems\": num_scenarios,\n",
    "                \"maxItems\": num_scenarios,\n",
    "            }\n",
    "        },\n",
    "        \"required\": [\"scenarios\"],\n",
    "        \"additionalProperties\": False,\n",
    "    }\n",
    "\n",
    "    # OpenRouter client (via OpenAI SDK)\n",
    "    try:\n",
    "        model = INPUT_GENERATION_MODEL  # noqa: F821 (defined elsewhere in your notebook)\n",
    "    except NameError:\n",
    "        model = \"openai/gpt-4.1-mini\"  # safe default if not set\n",
    "        warn(f\"INPUT_GENERATION_MODEL not defined; using default: {model}\")\n",
    "\n",
    "    step(f\"Calling OpenRouter model: {_C.BOLD}{model}{_C.RESET} …\")\n",
    "    client_openai = openai.OpenAI(\n",
    "        api_key=os.getenv(\"OPENROUTER_API_KEY\"),\n",
    "        base_url=\"https://openrouter.ai/api/v1\",\n",
    "    )\n",
    "\n",
    "    t1 = time.perf_counter()\n",
    "    response = client_openai.chat.completions.create(\n",
    "        model=model,\n",
    "        messages=[{\"role\": \"user\", \"content\": prompt}],\n",
    "        max_completion_tokens=8000,\n",
    "        response_format={\n",
    "            \"type\": \"json_schema\",\n",
    "            \"json_schema\": {\"name\": \"scenario_list\", \"schema\": response_schema},\n",
    "        },\n",
    "    )\n",
    "    dt = time.perf_counter() - t1\n",
    "    ok(f\"Model responded in {dt:.2f}s.\")\n",
    "\n",
    "    content = response.choices[0].message.content\n",
    "    info(f\"Raw content length: {len(content)} chars.\")\n",
    "    # Parse JSON\n",
    "    try:\n",
    "        result = json.loads(content)\n",
    "    except Exception as e:\n",
    "        err(\"Failed to parse JSON from model response.\")\n",
    "        dim(f\"   Exception: {e}\")\n",
    "        dim(\"   First 500 chars of response content:\")\n",
    "        dim(content[:500])\n",
    "        raise\n",
    "\n",
    "    # Extract scenarios\n",
    "    if \"scenarios\" in result:\n",
    "        scenarios = result[\"scenarios\"]\n",
    "    else:\n",
    "        scenarios = result if isinstance(result, list) else list(result.values())[0]\n",
    "\n",
    "    # Validate count\n",
    "    if len(scenarios) != num_scenarios:\n",
    "        err(f\"Expected {num_scenarios} scenarios, got {len(scenarios)}.\")\n",
    "        raise ValueError(f\"Expected {num_scenarios} scenarios, got {len(scenarios)}\")\n",
    "\n",
    "    ok(f\"Parsed {len(scenarios)} scenario(s) successfully.\")\n",
    "    preview_scenarios(scenarios, n=min(5, num_scenarios))\n",
    "    return scenarios\n",
    "\n",
    "\n",
    "# ---------- run generation w/ attempts ----------\n",
    "try:\n",
    "    expected_total = TRAINING_CONFIG[\"num_training_inputs\"] + NUM_TEST_INPUTS  # noqa: F821\n",
    "except NameError:\n",
    "    err(\"TRAINING_CONFIG/NUM_TEST_INPUTS not defined in this notebook.\")\n",
    "    raise\n",
    "\n",
    "info(f\"Target total scenarios: {expected_total}\")\n",
    "max_attempts = 10\n",
    "scenarios = None\n",
    "\n",
    "for attempt in range(1, max_attempts + 1):\n",
    "    step(f\"Attempt {attempt}/{max_attempts} …\")\n",
    "    t_attempt = time.perf_counter()\n",
    "    try:\n",
    "        scenarios = await generate_scenarios(num_scenarios=expected_total)\n",
    "        ok(f\"Attempt {attempt} succeeded in {time.perf_counter() - t_attempt:.2f}s.\")\n",
    "        break\n",
    "    except Exception as e:\n",
    "        warn(f\"Attempt {attempt} failed: {e}\")\n",
    "        if attempt < max_attempts:\n",
    "            time.sleep(min(1.5 * attempt, 6.0))\n",
    "        else:\n",
    "            err(\"All attempts exhausted.\")\n",
    "            raise\n",
    "\n",
    "# ---------- post-process & reporting ----------\n",
    "print()  # spacing\n",
    "ok(f\"Generated {len(scenarios)} scenarios total.\")\n",
    "info(\"Difficulty distribution:\")\n",
    "diff_counts = Counter(s[\"difficulty\"] for s in scenarios)\n",
    "for d in range(1, 6):\n",
    "    cnt = diff_counts.get(d, 0)\n",
    "    bar = \"█\" * min(cnt, 30)\n",
    "    dim(f\"   {d}/5: {cnt:3d}  {bar}\")\n",
    "\n",
    "print()\n",
    "step(\"Shuffling scenarios and splitting into train/val …\")\n",
    "random.shuffle(scenarios)\n",
    "\n",
    "train_n = TRAINING_CONFIG[\"num_training_inputs\"]  # noqa: F821\n",
    "raw_train_scenarios = scenarios[:train_n]\n",
    "raw_val_scenarios = scenarios[train_n:]\n",
    "\n",
    "ok(f\"Train: {len(raw_train_scenarios)} | Val: {len(raw_val_scenarios)}\")\n",
    "\n",
    "info(\"Sample (train) preview:\")\n",
    "preview_scenarios(raw_train_scenarios, n=min(5, len(raw_train_scenarios)))\n",
    "\n",
    "info(\"Sample (val) preview:\")\n",
    "preview_scenarios(raw_val_scenarios, n=min(5, len(raw_val_scenarios)))\n",
    "\n",
    "print()\n",
    "ok(\"Done.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "FET-_U0IPjq8"
   },
   "outputs": [],
   "source": [
    "# @title Run this cell to train your model!\n",
    "\n",
    "import os\n",
    "import random\n",
    "from dataclasses import dataclass\n",
    "\n",
    "import weave\n",
    "from dotenv import load_dotenv\n",
    "from openai import AsyncOpenAI\n",
    "\n",
    "import art\n",
    "from art.local import LocalBackend\n",
    "from art.rewards import ruler_score_group\n",
    "from art.utils import iterate_dataset\n",
    "\n",
    "load_dotenv()\n",
    "\n",
    "# Optional\n",
    "if WANDB_API_KEY:\n",
    "    os.environ[\"WANDB_API_KEY\"] = WANDB_API_KEY\n",
    "    weave.init(PROJECT_NAME)\n",
    "else:\n",
    "    print(\"WANDB_API_KEY is not set. We'll skip logging metrics to Weights & Biases.\")\n",
    "\n",
    "random.seed(42)\n",
    "\n",
    "# Declare the model\n",
    "model = art.TrainableModel(\n",
    "    name=MODEL_NAME,\n",
    "    project=PROJECT_NAME,\n",
    "    base_model=BASE_MODEL,\n",
    ")\n",
    "\n",
    "# To run on a T4, we need to override some config defaults.\n",
    "model._internal_config = art.dev.InternalModelConfig(\n",
    "    init_args=art.dev.InitArgs(\n",
    "        max_seq_length=MAX_SEQ_LENGTH,\n",
    "    ),\n",
    "    engine_args=art.dev.EngineArgs(\n",
    "        enforce_eager=True,\n",
    "        gpu_memory_utilization=GPU_MEMORY_UTILIZATION,\n",
    "    ),\n",
    ")\n",
    "\n",
    "# Initialize the server\n",
    "backend = LocalBackend(\n",
    "    in_process=True,\n",
    "    path=\"./.art\",\n",
    ")\n",
    "\n",
    "# Register the model with the local Backend\n",
    "await model.register(backend)\n",
    "\n",
    "print(\"Model created!\")\n",
    "print(\"Base model:\", BASE_MODEL)\n",
    "print(\"Model name:\", MODEL_NAME)\n",
    "print(\"Project name:\", PROJECT_NAME)\n",
    "\n",
    "\n",
    "def get_content_text(result) -> str:\n",
    "    # Extract text content from tool call result per MCP content schema\n",
    "    if isinstance(result, str):\n",
    "        return result\n",
    "    if hasattr(result, \"content\") and result.content:\n",
    "        out = \"\"\n",
    "        for item in result.content:\n",
    "            if isinstance(item, types.TextContent):\n",
    "                out += item.text\n",
    "            else:\n",
    "                out += str(item)\n",
    "        return out\n",
    "    if hasattr(result, \"structured\") and result.structured is not None:\n",
    "        try:\n",
    "            return json.dumps(result.structured)\n",
    "        except Exception:\n",
    "            return str(result.structured)\n",
    "    return str(result)\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class McpScenario:\n",
    "    \"\"\"A scenario for MCP agent evaluation against a remote Smithery server.\"\"\"\n",
    "\n",
    "    task_description: str\n",
    "    max_turns: int = MAX_TURNS\n",
    "\n",
    "\n",
    "@weave.op()\n",
    "async def rollout(\n",
    "    model: art.Model,\n",
    "    scenario: McpScenario,\n",
    "    debug: bool = False,\n",
    ") -> art.Trajectory:\n",
    "    \"\"\"Run an MCP agent rollout against the remote Smithery MCP server.\"\"\"\n",
    "    traj = art.Trajectory(\n",
    "        messages_and_choices=[],\n",
    "        reward=0,\n",
    "        metadata={\"task\": scenario.task_description},\n",
    "        metrics={\n",
    "            \"task_completed\": False,\n",
    "            \"success\": False,\n",
    "            \"ran_out_of_turns\": False,\n",
    "        },\n",
    "        scenario=scenario,\n",
    "    )\n",
    "\n",
    "    # Discover available tools from the remote server\n",
    "    tools_result, _resources_result = await list_tools_and_resources()\n",
    "    tool_names = [t.name for t in tools_result.tools]\n",
    "    log(\"rollout: discovered tools\", count=len(tool_names), names=tool_names)\n",
    "\n",
    "    # Convert to OpenAI tool format\n",
    "    tool_schemas = []\n",
    "    for tool in tools_result.tools:\n",
    "        tool_schema = {\n",
    "            \"type\": \"function\",\n",
    "            \"function\": {\n",
    "                \"name\": tool.name,\n",
    "                \"description\": tool.description or f\"MCP tool: {tool.name}\",\n",
    "                \"parameters\": tool.inputSchema or {\"type\": \"object\", \"properties\": {}},\n",
    "            },\n",
    "        }\n",
    "        tool_schemas.append(tool_schema)\n",
    "\n",
    "    # Add completion tool schema\n",
    "    tool_schemas.append(\n",
    "        {\n",
    "            \"type\": \"function\",\n",
    "            \"function\": {\n",
    "                \"name\": \"complete_task\",\n",
    "                \"description\": \"Complete the task with a summary\",\n",
    "                \"parameters\": {\n",
    "                    \"type\": \"object\",\n",
    "                    \"properties\": {\n",
    "                        \"summary\": {\n",
    "                            \"type\": \"string\",\n",
    "                            \"description\": \"Summary of accomplishments\",\n",
    "                        }\n",
    "                    },\n",
    "                    \"required\": [\"summary\"],\n",
    "                },\n",
    "            },\n",
    "        }\n",
    "    )\n",
    "\n",
    "    traj.tools = tool_schemas\n",
    "\n",
    "    # Initialize conversation\n",
    "    system_prompt = (\n",
    "        f\"You are an MCP (Model Context Protocol) agent.\\n\\n\"\n",
    "        f\"Use MCP tools through the server to complete your task.\\n\\n\"\n",
    "        f\"When you believe you have completed the task, call the 'complete_task' function with a summary of what you accomplished. \"\n",
    "        f\"You have a total of {scenario.max_turns} turns.\"\n",
    "        # NOTE: removing 'Only use tool calls, do not write any content.' — some models\n",
    "        # will freeze if they think plain text is disallowed. Let them output thoughts but\n",
    "        # we only process tool calls below.\n",
    "    )\n",
    "\n",
    "    traj.messages_and_choices = [\n",
    "        {\"role\": \"system\", \"content\": system_prompt},\n",
    "        {\n",
    "            \"role\": \"user\",\n",
    "            \"content\": f\"Please complete this task: {scenario.task_description}\",\n",
    "        },\n",
    "    ]\n",
    "\n",
    "    num_turns = 0\n",
    "    task_completed = False\n",
    "\n",
    "    # Main interaction loop\n",
    "    while num_turns < scenario.max_turns and not task_completed:\n",
    "        num_turns += 1\n",
    "\n",
    "        try:\n",
    "            # === Log request ===\n",
    "            last_user = next(\n",
    "                (m for m in reversed(traj.messages()) if m[\"role\"] == \"user\"), None\n",
    "            )\n",
    "            log(\n",
    "                \"LLM request\",\n",
    "                step=num_turns,\n",
    "                model=(model.inference_model_name or model.name),\n",
    "                tools=len(tool_schemas),\n",
    "                last_user=(last_user[\"content\"][:160] + \"...\" if last_user else None),\n",
    "            )\n",
    "\n",
    "            # Get LLM response\n",
    "            async with traj.track_duration(\"llm_completion\"):\n",
    "                openai_client = AsyncOpenAI(\n",
    "                    api_key=model.inference_api_key,\n",
    "                    base_url=model.inference_base_url,\n",
    "                )\n",
    "\n",
    "                # We also log the request body (without huge params)\n",
    "                req_preview = {\n",
    "                    \"model\": model.inference_model_name\n",
    "                    if model.inference_model_name\n",
    "                    else model.name,\n",
    "                    \"messages_len\": len(traj.messages()),\n",
    "                    \"tools_len\": len(tool_schemas),\n",
    "                }\n",
    "                log_json(\"LLM request (preview)\", req_preview)\n",
    "\n",
    "                response = await openai_client.chat.completions.create(\n",
    "                    model=model.inference_model_name\n",
    "                    if model.inference_model_name\n",
    "                    else model.name,\n",
    "                    messages=traj.messages(),\n",
    "                    tools=tool_schemas,\n",
    "                    max_completion_tokens=8000,\n",
    "                )\n",
    "\n",
    "            # === Log response ===\n",
    "            choice = response.choices[0]\n",
    "\n",
    "            finish_reason = getattr(choice, \"finish_reason\", None)\n",
    "            msg = choice.message\n",
    "            has_tools = bool(getattr(msg, \"tool_calls\", None))\n",
    "            content_preview = (\n",
    "                (msg.content[:200] + \"...\")\n",
    "                if isinstance(msg.content, str) and msg.content\n",
    "                else str(msg.content)[:200]\n",
    "            )\n",
    "            log(\n",
    "                \"LLM response parsed\",\n",
    "                finish_reason=finish_reason,\n",
    "                has_tool_calls=has_tools,\n",
    "                content_preview=content_preview,\n",
    "            )\n",
    "\n",
    "            traj.messages_and_choices.append(choice)\n",
    "\n",
    "            # Handle tool calls\n",
    "            if msg.tool_calls:\n",
    "                for tool_call in msg.tool_calls:\n",
    "                    try:\n",
    "                        log(\n",
    "                            \"Tool call received\",\n",
    "                            name=tool_call.function.name,\n",
    "                            raw_args=tool_call.function.arguments,\n",
    "                        )\n",
    "                        tool_args = json.loads(tool_call.function.arguments or \"{}\")\n",
    "\n",
    "                        if tool_call.function.name == \"complete_task\":\n",
    "                            traj.metrics[\"task_completed\"] = True\n",
    "                            task_completed = True\n",
    "                            traj.logs.append(\n",
    "                                f\"Task completion attempted with summary: {tool_args.get('summary', '')}\"\n",
    "                            )\n",
    "                            # We still append a tool message for completeness\n",
    "                            traj.messages_and_choices.append(\n",
    "                                {\n",
    "                                    \"role\": \"tool\",\n",
    "                                    \"tool_call_id\": tool_call.id,\n",
    "                                    \"content\": \"Task marked complete.\",\n",
    "                                }\n",
    "                            )\n",
    "                        else:\n",
    "                            # 🔧 Call MCP tool through remote Smithery session\n",
    "                            result = await call_mcp_tool(\n",
    "                                tool_call.function.name, tool_args\n",
    "                            )\n",
    "\n",
    "                            content_text = get_content_text(result)\n",
    "                            log(\n",
    "                                \"Tool result\",\n",
    "                                name=tool_call.function.name,\n",
    "                                len=len(content_text),\n",
    "                            )\n",
    "\n",
    "                            if len(content_text) > 20000:\n",
    "                                # print(\n",
    "                                #     f\"Tool call result for {tool_call.function.name} is too long: {len(content_text)}\"\n",
    "                                # )\n",
    "                                # print(f\"Args: {tool_args}\")\n",
    "                                # print(content_text[:1000])\n",
    "                                # print(content_text[-1000:])\n",
    "                                raise Exception(\n",
    "                                    f\"Tool call result for {tool_call.function.name} is too long: {len(content_text)}\"\n",
    "                                )\n",
    "\n",
    "                            # Add tool response\n",
    "                            traj.messages_and_choices.append(\n",
    "                                {\n",
    "                                    \"role\": \"tool\",\n",
    "                                    \"tool_call_id\": tool_call.id,\n",
    "                                    \"content\": content_text,\n",
    "                                }\n",
    "                            )\n",
    "\n",
    "                    except Exception as e:\n",
    "                        traceback.print_exc()\n",
    "                        traj.logs.append(f\"Tool call error: {e}\")\n",
    "\n",
    "                        # Add error response\n",
    "                        traj.messages_and_choices.append(\n",
    "                            {\n",
    "                                \"role\": \"tool\",\n",
    "                                \"tool_call_id\": tool_call.id,\n",
    "                                \"content\": f\"Error: {str(e)}\",\n",
    "                            }\n",
    "                        )\n",
    "            else:\n",
    "                # No tool calls — log and continue (RULER will likely give 0)\n",
    "                log(\n",
    "                    \"LLM returned no tool_calls; skipping tool execution\",\n",
    "                    turn=num_turns,\n",
    "                )\n",
    "                # You can consider breaking here or letting it try another turn\n",
    "                # break\n",
    "\n",
    "        except Exception as e:\n",
    "            traceback.print_exc()\n",
    "            traj.logs.append(f\"Error in turn {num_turns}: {e}\")\n",
    "            break\n",
    "\n",
    "    if not task_completed and num_turns == scenario.max_turns:\n",
    "        traj.metrics[\"ran_out_of_turns\"] = True\n",
    "\n",
    "    traj.metrics[\"num_turns\"] = num_turns\n",
    "\n",
    "    return traj.finish()\n",
    "\n",
    "\n",
    "# =============== Training code ===============\n",
    "\n",
    "print(\n",
    "    f\"Using config: max_turns={MAX_TURNS}, rollouts_per_group={TRAINING_CONFIG['rollouts_per_group']}, \"\n",
    "    f\"groups_per_step={TRAINING_CONFIG['groups_per_step']}, num_epochs={TRAINING_CONFIG['num_epochs']}, \"\n",
    "    f\"learning_rate={TRAINING_CONFIG['learning_rate']}\"\n",
    ")\n",
    "\n",
    "await model.register(backend)\n",
    "\n",
    "train_scenarios = [\n",
    "    McpScenario(\n",
    "        task_description=scenario[\"task\"],\n",
    "        max_turns=MAX_TURNS,\n",
    "    )\n",
    "    for scenario in raw_train_scenarios\n",
    "]\n",
    "\n",
    "# Create dataset iterator using raw scenarios\n",
    "train_iterator = iterate_dataset(\n",
    "    train_scenarios,\n",
    "    groups_per_step=TRAINING_CONFIG[\"groups_per_step\"],\n",
    "    num_epochs=TRAINING_CONFIG[\"num_epochs\"],\n",
    "    initial_step=await model.get_step(),  # Resume from checkpoint\n",
    ")\n",
    "\n",
    "# Main training loop using iterate_dataset\n",
    "for batch in train_iterator:\n",
    "    print(\"Gathering trajectory groups with RULER scoring...\")\n",
    "\n",
    "    # Use gather_trajectory_groups with ruler_score_group\n",
    "    groups = await art.gather_trajectory_groups(\n",
    "        (\n",
    "            art.TrajectoryGroup(\n",
    "                rollout(model, scenario, False)\n",
    "                for _ in range(TRAINING_CONFIG[\"rollouts_per_group\"])\n",
    "            )\n",
    "            for scenario in batch.items\n",
    "        ),\n",
    "        pbar_desc=f\"train gather step {batch.step}\",\n",
    "    )\n",
    "\n",
    "    scored_groups = []\n",
    "    for group in groups:\n",
    "        # Use RULER to assign relative scores to each trajectory\n",
    "        judged_group = await ruler_score_group(\n",
    "            group, judge_model=RULER_MODEL, debug=True, swallow_exceptions=True\n",
    "        )\n",
    "        scored_groups.append(judged_group)\n",
    "\n",
    "    print(\"starting train\")\n",
    "    await model.train(\n",
    "        scored_groups,\n",
    "        config=art.TrainConfig(learning_rate=TRAINING_CONFIG[\"learning_rate\"]),\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "YRO9ndqo5ky4"
   },
   "outputs": [],
   "source": [
    "# @title Test Your Model!\n",
    "\n",
    "# Generate test inputs\n",
    "print(\"Generating test inputs...\")\n",
    "val_scenarios = [\n",
    "    McpScenario(\n",
    "        task_description=scenario[\"task\"],\n",
    "        max_turns=MAX_TURNS,\n",
    "    )\n",
    "    for scenario in raw_val_scenarios\n",
    "]\n",
    "\n",
    "print(f\"\\n🧪 Testing the trained model on {len(val_scenarios)} new inputs:\\n\")\n",
    "print(\"=\" * 80)\n",
    "\n",
    "for i, scenario in enumerate(val_scenarios):\n",
    "    print(f\"\\nTest {i + 1}:\")\n",
    "    print(f\"Input: {scenario.task_description}\")\n",
    "\n",
    "    # Run the model\n",
    "    result_trajectory = await rollout(model, scenario)\n",
    "\n",
    "    # Extract the model's response\n",
    "    messages = result_trajectory.messages()\n",
    "    model_response = messages[-1][\"content\"] if messages else \"No response\"\n",
    "\n",
    "    print(f\"Model output: {model_response}\")\n",
    "    print(\"-\" * 80)\n",
    "\n",
    "print(\"\\n🎉 Testing completed!\")\n",
    "print(\n",
    "    f\"\\nYour model '{MODEL_NAME}' has been trained to use the Smithery MCP server at:\"\n",
    ")\n",
    "print(SMITHERY_MCP_URL)\n",
    "print(\"\\nTo use this model in production:\")\n",
    "print(\"1. The model checkpoint is saved in ./.art/\")\n",
    "print(\"2. You can load it using the vLLM library\")\n",
    "print(\n",
    "    \"3. Or continue training with more examples by adjusting the configuration at the top\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "utI-VYM8s5lo"
   },
   "outputs": [],
   "source": [
    "# @title Upload to Hugging Face 🤗\n",
    "\n",
    "import torch\n",
    "from unsloth import FastLanguageModel\n",
    "\n",
    "lora_model_path = (\n",
    "    f\".art/{model.project}/models/{model.name}/{await model.get_step():04d}\"\n",
    ")\n",
    "\n",
    "peft_model, tokenizer = FastLanguageModel.from_pretrained(\n",
    "    model_name=lora_model_path,\n",
    "    max_seq_length=16384,\n",
    "    dtype=torch.bfloat16,\n",
    "    load_in_4bit=True,\n",
    ")\n",
    "\n",
    "if False:  # Change to True to upload finetune\n",
    "    peft_model.push_to_hub_merged(f\"HF_ACCOUNT/{model.name}\", tokenizer, token=\"hf_...\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "FuevYgXT-I1h"
   },
   "source": [
    "### Next Steps\n",
    "\n",
    "Congratulations! You've successfully trained a custom model for your task using only:\n",
    "- A pre-built MCP server\n",
    "- Example inputs (no outputs needed!)\n",
    "- RULER's automatic evaluation\n",
    "\n",
    "Here are some ways to improve results:\n",
    "\n",
    "1. **More diverse inputs**: Generate more varied input examples\n",
    "2. **Longer training**: Increase the number of training steps\n",
    "3. **More comparisons**: Increase `rollouts_per_group` for better RULER comparisons\n",
    "4. **MCP server refinement**: Add better tools and resources to the server\n",
    "5. **Hyperparameter tuning**: Adjust learning rate, batch size, etc.\n",
    "\n",
    "Remember: RULER learns what \"good\" means from your MCP server alone - no labeled data required!\n",
    "\n",
    "For more advanced use cases, check out the [ART documentation](https://art.openpipe.ai)."
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "gpuType": "T4",
   "provenance": [
    {
     "file_id": "https://github.com/openpipe/art/blob/main/examples/mcp-rl/mcp-rl.ipynb",
     "timestamp": 1754988958411
    }
   ]
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
